#include "traj_utils/uniform_bspline.h"
#include <ros/ros.h>

namespace ego_planner {

/**
 * @brief De Boor算法是用于计算B样条曲线上的点的标准方法，
 *          通过对参数 u 所在的区间进行插值，以获取曲线上的点坐标
 *
 * @param u
 * @return Eigen::VectorXd
 */
Eigen::VectorXd UniformBspline::getPos(const double &u)
{

  double ub = min(max(u_(p_), u), u_(m_ - p_));

  // determine which [ui,ui+1] lay in
  int k = p_;
  while (true)
  {
    if (u_(k + 1) >= ub)
      break;
    ++k;
  }

  /* deBoor's alg */
  vector<Eigen::VectorXd> d;
  for (int i = 0; i <= p_; ++i)
  {
    d.push_back(control_points_.col(k - p_ + i));
    // cout << d[i].transpose() << endl;
  }

  for (int r = 1; r <= p_; ++r)
  {
    for (int i = p_; i >= r; --i)
    {
      double alpha = (ub - u_[i + k - p_]) / (u_[i + 1 + k - r] - u_[i + k - p_]);
      // cout << "alpha: " << alpha << endl;
      d[i] = (1 - alpha) * d[i - 1] + alpha * d[i];
    }
  }

  return d[p_];
}

Eigen::MatrixXd UniformBspline::getDerivativeControlPoints()
{
  // The derivative of a b-spline is also a b-spline, its order become p_-1
  // control point Qi = p_*(Pi+1-Pi)/(ui+p_+1-ui+1)
  Eigen::MatrixXd ctp(control_points_.rows(), control_points_.cols() - 1);
  for (int i = 0; i < ctp.cols(); ++i)
  {
    ctp.col(i) =
        p_ * (control_points_.col(i + 1) - control_points_.col(i)) / (u_(i + p_ + 1) - u_(i + 1));
  }
  return ctp;
}

UniformBspline UniformBspline::getDerivative()
{
  Eigen::MatrixXd ctp = getDerivativeControlPoints();
  UniformBspline derivative(ctp, p_ - 1, interval_);

  /* cut the first and last knot */
  Eigen::VectorXd knot(u_.rows() - 2);
  knot = u_.segment(1, u_.rows() - 2);
  derivative.setKnot(knot);

  return derivative;
}

void UniformBspline::setPhysicalLimits(const double &vel,
                                       const double &acc,
                                       const double &tolerance)
{
  limit_vel_ = vel;
  limit_acc_ = acc;
  limit_ratio_ = 1.1;
  feasibility_tolerance_ = tolerance;
}

/**
 * @param ratio the ratio between real vel(acc) and limit vel(acc), derive from formulation(14)
 * @param show  whether display the prompt message
 */
bool UniformBspline::checkFeasibility(double &ratio, bool show)
{
  bool fea = true;

  Eigen::MatrixXd P = control_points_;
  int dimension = control_points_.rows();

  /* check vel feasibility and insert points */
  double max_vel = -1.0;
  double enlarged_vel_lim = limit_vel_ * (1.0 + feasibility_tolerance_) + 1e-4;
  for (int i = 0; i < P.cols() - 1; ++i)
  {
    Eigen::VectorXd vel = p_ * (P.col(i + 1) - P.col(i)) / (u_(i + p_ + 1) - u_(i + 1));

    if (fabs(vel(0)) > enlarged_vel_lim || fabs(vel(1)) > enlarged_vel_lim ||
        fabs(vel(2)) > enlarged_vel_lim)
    {
      if (show)
        cout << "[Check]: Infeasible vel " << i << " :" << vel.transpose() << endl;
      fea = false;

      for (int j = 0; j < dimension; ++j) {
        max_vel = max(max_vel, fabs(vel(j)));
      }
    }
  }

  /* acc feasibility */
  double max_acc = -1.0;
  double enlarged_acc_lim = limit_acc_ * (1.0 + feasibility_tolerance_) + 1e-4;
  for (int i = 0; i < P.cols() - 2; ++i)
  {
    Eigen::VectorXd acc = p_ * (p_ - 1) *
                          ((P.col(i + 2) - P.col(i + 1)) / (u_(i + p_ + 2) - u_(i + 2)) -
                           (P.col(i + 1) - P.col(i)) / (u_(i + p_ + 1) - u_(i + 1))) /
                          (u_(i + p_ + 1) - u_(i + 2));

    if (fabs(acc(0)) > enlarged_acc_lim || fabs(acc(1)) > enlarged_acc_lim ||
        fabs(acc(2)) > enlarged_acc_lim)
    {
      if (show)
        cout << "[Check]: Infeasible acc " << i << " :" << acc.transpose() << endl;
      fea = false;

      for (int j = 0; j < dimension; ++j) {
        max_acc = max(max_acc, fabs(acc(j)));
      }
    }
  }

  ratio = max(max_vel / limit_vel_, sqrt(fabs(max_acc) / limit_acc_));

  return fea;
}

/**
 * 延长b-spline的 knots 的时间间隔
 * @note 这里直接从第7个开始改变，因为前6(num1+1)个点中，前三个是自由点，中间三个
 *       会对初始条件产生影响，所以不能随意改变
 */
void UniformBspline::lengthenTime(const double &ratio)
{
  cout << "p_ = " << p_ << endl;
  cout << "u_.rows() = " << u_.rows() << endl;
  cout << "m_ = " << m_ << endl;
  int num1 = 5;
  int num2 = getKnot().rows() - 1 - 5;

  double delta_t = (ratio - 1.0) * (u_(num2) - u_(num1));
  double t_inc = delta_t / double(num2 - num1);
  for (int i = num1 + 1; i <= num2; ++i)
    u_(i) += double(i - num1) * t_inc;
  for (int i = num2 + 1; i < u_.rows(); ++i)
    u_(i) += delta_t;
}

/**
 * @brief B-spline fitting with boundary vel&acc constraints and time interval
 *
 * @param dt: time interval
 * @param pts_init: initial points which need to be fitted (pass through)
 * @param terminal_constraint: boundary vel&acc
 * @param pts_ctrl: control points of B-spline after fitting
 *
 *  TODO 感觉边界条件给多了，导致出来的是一个列满秩矩阵，最后相当于是在用最小而乘求解
 *  TODO 这里使用的不是clamped b-spline，是均匀b-spline取了中间一段
 */
void UniformBspline::BsplineFitting(const double &dt,
                                    const vector<Eigen::Vector3d> &pts_init,
                                    const vector<Eigen::Vector3d> &terminal_constraint,
                                    Eigen::MatrixXd &pts_ctrl)
{
  if (dt <= 0)
  {
    cout << "[B-spline]:time step error." << endl;
    return;
  }

  if (pts_init.size() <= 3)
  {
    cout << "[B-spline]:point set have only " << pts_init.size() << " points." << endl;
    return;
  }

  if (terminal_constraint.size() != 4)
  {
    cout << "[B-spline]:derivatives error." << endl;
  }

  size_t K = pts_init.size();

  // write A
  Eigen::Vector3d prow(3), vrow(3), arow(3);
  prow << 1, 4, 1;
  vrow << -1, 0, 1;
  arow << 1, -2, 1;

  Eigen::MatrixXd A = Eigen::MatrixXd::Zero(K + 4, K + 2);

  for (size_t i = 0; i < K; ++i)
    A.block(i, i, 1, 3) = (1 / 6.0) * prow.transpose();

  A.block(K, 0, 1, 3) = (1 / 2.0 / dt) * vrow.transpose();
  A.block(K + 1, K - 1, 1, 3) = (1 / 2.0 / dt) * vrow.transpose();

  A.block(K + 2, 0, 1, 3) = (1 / dt / dt) * arow.transpose();
  A.block(K + 3, K - 1, 1, 3) = (1 / dt / dt) * arow.transpose();

  // cout << "A" << endl << A << endl << endl;

  // write b
  Eigen::VectorXd bx(K + 4), by(K + 4), bz(K + 4);
  for (size_t i = 0; i < K; ++i)
  {
    bx(i) = pts_init[i](0);
    by(i) = pts_init[i](1);
    bz(i) = pts_init[i](2);
  }

  for (size_t i = 0; i < 4; ++i)
  {
    bx(K + i) = terminal_constraint[i](0);
    by(K + i) = terminal_constraint[i](1);
    bz(K + i) = terminal_constraint[i](2);
  }

  // solve Ax = b
  Eigen::VectorXd px = A.colPivHouseholderQr().solve(bx);
  Eigen::VectorXd py = A.colPivHouseholderQr().solve(by);
  Eigen::VectorXd pz = A.colPivHouseholderQr().solve(bz);

  // convert to control pts
  pts_ctrl.resize(3, K + 2);
  pts_ctrl.row(0) = px.transpose();
  pts_ctrl.row(1) = py.transpose();
  pts_ctrl.row(2) = pz.transpose();

  // cout << "[B-spline]: parameterization ok." << endl;
}

double UniformBspline::getJerk()
{
  UniformBspline jerk_traj = getDerivative().getDerivative().getDerivative();

  Eigen::VectorXd times = jerk_traj.getKnot();
  Eigen::MatrixXd ctrl_pts = jerk_traj.getControlPoint();
  int dimension = ctrl_pts.rows();

  double jerk = 0.0;
  for (int i = 0; i < ctrl_pts.cols(); ++i)
  {
    for (int j = 0; j < dimension; ++j)
    {
      jerk += (times(i + 1) - times(i)) * ctrl_pts(j, i) * ctrl_pts(j, i);
    }
  }

  return jerk;
}

void UniformBspline::getMeanAndMaxVel(double &mean_v, double &max_v)
{
  UniformBspline vel = getDerivative();
  double tm, tmp;
  vel.getTimeSpan(tm, tmp);

  double max_vel = -1.0, mean_vel = 0.0;
  int num = 0;
  for (double t = tm; t <= tmp; t += 0.01)
  {
    Eigen::VectorXd vxd = vel.getPos(t);
    double vn = vxd.norm();

    mean_vel += vn;
    ++num;
    if (vn > max_vel)
    {
      max_vel = vn;
    }
  }

  mean_vel = mean_vel / double(num);
  mean_v = mean_vel;
  max_v = max_vel;
}

void UniformBspline::getMeanAndMaxAcc(double &mean_a, double &max_a)
{
  UniformBspline acc = getDerivative().getDerivative();
  double tm, tmp;
  acc.getTimeSpan(tm, tmp);

  double max_acc = -1.0, mean_acc = 0.0;
  int num = 0;
  for (double t = tm; t <= tmp; t += 0.01)
  {
    Eigen::VectorXd axd = acc.getPos(t);
    double an = axd.norm();

    mean_acc += an;
    ++num;
    if (an > max_acc)
    {
      max_acc = an;
    }
  }

  mean_acc = mean_acc / double(num);
  mean_a = mean_acc;
  max_a = max_acc;
}
} // namespace ego_planner
